# @Version : 1.0
# @Author  : wxz
# @File    : chat_service.py
# @Time    : 2025/6/18 23:17
# @Describe: 描述...
import json
from datetime import datetime
from typing import Mapping, Any, Optional, List

from autogen_agentchat.state import AssistantAgentState
from autogen_core.model_context import ChatCompletionContextState
from fastapi import Depends
from sqlalchemy.orm import Session

from apps.chat.agent.chat_agent import ChatAgent
from apps.chat.dao.chat_message_dao import ChatMessageDao
from apps.chat.dao.user_chat_dao import UserChatDao
from apps.chat.data_model.chat_message_model import ChatMessageModel
from apps.chat.data_model.user_chat_model import UserChatModel
from apps.chat.exception.chat_exception import UserChatNotHave
from apps.chat.schemas.chat_message_dto import ChatMessageDto
from apps.chat.schemas.chat_response import ChatResponse
from utils.id_util import IdUtil


class ChatService:

	def __init__(self,
				 chat_agent: ChatAgent = Depends(),
				 user_chat_dao: UserChatDao = Depends(),
				 chat_message_dao: ChatMessageDao = Depends(),
				 ):
		self.chat_agent = chat_agent
		self.user_chat_dao = user_chat_dao
		self.chat_message_dao = chat_message_dao

	def _save_chat_message(self, llm_context_messages: List[Mapping[str, Any]], user_id: str, user_chat_id: str):

		# 保存用户会话信息
		chat_message_lst = []
		for msg in llm_context_messages:
			chat_message = ChatMessageModel(
				id=IdUtil.get_uuid(),
				user_id=user_id,
				user_chat_id=user_chat_id,
				message=json.dumps(msg, ensure_ascii=False),
				serial_number=0,
			)
			chat_message_lst.append(chat_message)

		self.chat_message_dao.save_all(chat_message_lst)

	async def chat_messages(self, message: ChatMessageDto, user_id: str)-> ChatResponse:

		# Agent 会话状态
		agent_state: Optional[AssistantAgentState] = None

		# 判断是否需要加载会话状态
		if message.conversation_id is None or message.conversation_id == "":
			# 首次创建的会话
			# 1、获取一个 conversation_id
			conversation_id = IdUtil.get_uuid()
			# 2、
			chat_agent_response = await self.chat_agent.run_chat_agent(message.query, agent_state)

			# 保存用户会话
			user_chat = UserChatModel(
				id=IdUtil.get_uuid(),
				user_id=user_id,
				conversation_id=conversation_id,
				topic="用户聊天",
			)

			self.user_chat_dao.save_user_chat(user_chat)

			# 保存用户会话信息
			self._save_chat_message(chat_agent_response.llm_context_messages, user_id, user_chat.id)

			# 提交事务
			self.user_chat_dao.db_session.commit()
			self.chat_message_dao.db_session.commit()

			chat_response = ChatResponse(
				chat_content=chat_agent_response.content,
				conversation_id=user_chat.conversation_id,
			)

			return chat_response

		# 如果 conversation_id 会话ID不能为空则查询用户的历史聊天记录
		# 1、查询用户的聊天会话 user_chat
		user_chat_model = self.user_chat_dao.find_user_chat(user_id, message.conversation_id)
		# 会话聊天信息不存在
		if user_chat_model is None:
			raise UserChatNotHave

		# 查询 用户的聊天会话信息
		chat_message_lst = self.chat_message_dao.find_by_user_chat_id(user_id, user_chat_model.id)
		if chat_message_lst is None:
			raise UserChatNotHave

		# 取出 会话聊天内容JSON
		msg_list = []
		for chat_message in chat_message_lst:
			msg_dict = json.loads(chat_message.message)
			msg_list.append(msg_dict)

		# 创建Agent的会话状态
		messages: Optional[Mapping[str, Any]] = ChatCompletionContextState(messages=msg_list).model_dump()
		agent_state = AssistantAgentState(llm_context = messages)

		chat_agent_response = await self.chat_agent.run_chat_agent(message.query, agent_state)

		# 保存用户会话信息
		# 获取最后两条message
		msg_lst = chat_agent_response.llm_context_messages
		tmp_lst = [msg_lst[len(msg_lst) - 2], msg_lst[len(msg_lst) - 1]]

		self._save_chat_message(tmp_lst, user_id, user_chat_model.id)
		self.chat_message_dao.db_session.commit()

		chat_response = ChatResponse(
			chat_content=chat_agent_response.content,
			conversation_id=user_chat_model.conversation_id,
		)

		return chat_response

